from core.layer import Layer
from implement.layers.basic.batchnorm import BatchNorm
from implement.layers.conv.conv2d import Conv2d
from utils.functions_collect import relu


class BottleneckA(Layer):
    """A bottleneck layer that reduces the resolution of the feature map.
    Args:
        in_channels (int): Number of channels of input arrays.
        mid_channels (int): Number of channels of intermediate arrays.
        out_channels (int): Number of channels of output arrays.
        stride (int or tuple of ints): Stride of filter application.
        downsample_fb (bool): If this argument is specified as ``False``,
            it performs downsampling by placing stride 2
            on the 1x1 convolutional layers (the original MSRA ResNet).
            If this argument is specified as ``True``, it performs downsampling
            by placing stride 2 on the 3x3 convolutional layers
            (Facebook ResNet).
    """

    def __init__(self, in_channels, mid_channels, out_channels,
                 stride=2, downsample_fb=False):
        super().__init__()
        # In the original MSRA ResNet, stride=2 is on 1x1 convolution.
        # In Facebook ResNet, stride=2 is on 3x3 convolution.
        stride_1x1, stride_3x3 = (1, stride) if downsample_fb else (stride, 1)

        self.conv1 = Conv2d(mid_channels, 1, stride_1x1, 0,
                            nobias=True)
        self.bn1 = BatchNorm()
        self.conv2 = Conv2d(mid_channels, 3, stride_3x3, 1,
                            nobias=True)
        self.bn2 = BatchNorm()
        self.conv3 = Conv2d(out_channels, 1, 1, 0, nobias=True)
        self.bn3 = BatchNorm()
        self.conv4 = Conv2d(out_channels, 1, stride, 0,
                            nobias=True)
        self.bn4 = BatchNorm()

    def forward(self, x):
        h1 = relu(self.bn1(self.conv1(x)))
        h1 = relu(self.bn2(self.conv2(h1)))
        h1 = self.bn3(self.conv3(h1))
        h2 = self.bn4(self.conv4(x))
        return relu(h1 + h2)
